import torch
import numpy as np
import time
from src.optimizers.factory import OptimizerFactory
from src.problems.coco_benchmark import CocoBenchmark
from src.models.rbfn import RBFN
from src.models.mlp import MLP
from src.models.gpr import GPR

def run_optimization(model, optimizer_type, test_function, dim, bounds, num_samples):
    """运行单次优化实验"""
    # 准备训练数据
    X = torch.FloatTensor(np.random.uniform(bounds[0], bounds[1], (num_samples, dim)))
    y = torch.FloatTensor([[test_function(x.numpy())] for x in X])
    
    # 训练模型
    start_time = time.time()
    model.fit(
        x=X,
        y=y,
        epochs=100,
        batch_size=32,
        learning_rate=0.01,
        validation_split=0.2
    )
    training_time = time.time() - start_time
    
    # 创建优化器
    optimizer = OptimizerFactory.create_optimizer(
        optimizer_type=optimizer_type,
        dim=dim,
        pop_size=50,
        max_iter=1000
    )
    
    # 运行优化
    start_time = time.time()
    best_solution, best_fitness = optimizer.optimize(
        objective_func=test_function,
        bounds=bounds,
        model=model
    )
    optimization_time = time.time() - start_time
    
    return {
        'best_fitness': best_fitness,
        'training_time': training_time,
        'optimization_time': optimization_time,
        'total_time': training_time + optimization_time
    }

def main():
    # 实验参数
    dim = 30
    bounds = (-5, 5)
    num_samples = 1000
    num_runs = 5  # 每个配置运行5次
    
    # 测试函数
    test_functions = {
        'Rastrigin': CocoBenchmark.rastrigin,
        'Sphere': CocoBenchmark.sphere,
        'Rosenbrock': CocoBenchmark.rosenbrock
    }
    
    # 模型配置
    models = {
        'RBFN': lambda: RBFN(
            input_dim=dim,
            output_dim=1,
            num_centers=50,
            sigma=1.0,
            use_kmeans_init=True,
            normalize_data=True
        ),
        'MLP': lambda: MLP(
            input_dim=dim,
            output_dim=1,
            hidden_dims=[128, 64],
            dropout_rate=0.1,
            activation='relu',
            batch_norm=True,
            normalize_data=True
        ),
        'GPR': lambda: GPR(
            input_dim=dim,
            output_dim=1,
            kernel_type='rbf',
            lengthscale_prior=1.0,
            noise_prior=0.1,
            normalize_data=True
        )
    }
    
    # 优化器类型
    optimizers = ['de', 'pso', 'ga']
    
    # 运行实验
    results = {}
    
    for func_name, test_function in test_functions.items():
        print(f"\n测试函数: {func_name}")
        print("-" * 50)
        
        for model_name, model_creator in models.items():
            for optimizer_type in optimizers:
                print(f"\n模型: {model_name}, 优化器: {optimizer_type}")
                
                # 运行多次实验
                run_results = []
                for run in range(num_runs):
                    print(f"运行 {run + 1}/{num_runs}")
                    model = model_creator()
                    result = run_optimization(
                        model, optimizer_type, test_function,
                        dim, bounds, num_samples
                    )
                    run_results.append(result)
                
                # 计算平均结果
                avg_fitness = np.mean([r['best_fitness'] for r in run_results])
                avg_training_time = np.mean([r['training_time'] for r in run_results])
                avg_optimization_time = np.mean([r['optimization_time'] for r in run_results])
                avg_total_time = np.mean([r['total_time'] for r in run_results])
                
                # 存储结果
                key = f"{func_name}_{model_name}_{optimizer_type}"
                results[key] = {
                    'avg_fitness': avg_fitness,
                    'avg_training_time': avg_training_time,
                    'avg_optimization_time': avg_optimization_time,
                    'avg_total_time': avg_total_time
                }
                
                # 打印结果
                print(f"平均最优适应度值: {avg_fitness:.6f}")
                print(f"平均训练时间: {avg_training_time:.2f}秒")
                print(f"平均优化时间: {avg_optimization_time:.2f}秒")
                print(f"平均总时间: {avg_total_time:.2f}秒")
    
    # 保存结果
    import json
    with open('benchmark_results.json', 'w') as f:
        json.dump(results, f, indent=4)
    
    print("\n实验结果已保存到 benchmark_results.json")

if __name__ == "__main__":
    main() 